
import torch

from collections import defaultdict


global_device = None
global_loss_dict = None
global_update_functions = None

class LossAccumulator(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.loss = torch.tensor(0, device=global_device)
        self.n = 0

    def add(self, loss):
        self.loss = self.loss + loss
        self.n += 1

    def avg(self):
        return self.loss / self.n if self.n else self.loss


def init_train_utils(device, **kwargs):
    global global_device
    global_device = device
    reset_global_loss_dict()
    global global_update_functions
    global_update_functions = []

def accum_loss(key: str, loss: torch.Tensor):
    global_loss_dict[key].add(loss)

def get_global_loss_dict():
    return global_loss_dict

def reset_global_loss_dict():
    global global_loss_dict
    global_loss_dict = defaultdict(LossAccumulator)

def register_update_function(func):
    if global_update_functions is not None:
        global_update_functions.append(func)

def run_global_update_functions():
    for update_func in global_update_functions:
        update_func()
